import paddle.fluid as fluid
import paddle
import os
import numpy as np
from PIL import Image
import cv2
import math
import add_mosaic.reader as reader
from add_mosaic.pyramidbox import PyramidBox


# os.environ['FLAGS_fraction_of_gpu_memory_to_use']='0.01'

class FaceDetection(object):

    def __init__(self, confs_threshold=0.7):
        self.work_dir = os.path.dirname(os.path.abspath(__file__))
        self.confs_threshold = confs_threshold
        self.model_dir = os.path.join(self.work_dir, '../model_data/face/')
        self.image_shape = [3, 1024, 1024]
        self.use_pyramidbox = True

        if paddle.device.is_compiled_with_cuda():
            place = fluid.CUDAPlace(0)
        else:
            place = fluid.CPUPlace()

        exe = fluid.Executor(place)
        main_program = fluid.Program()
        startup_program = fluid.Program()
        with fluid.program_guard(main_program, startup_program):
            network = PyramidBox(
                data_shape=self.image_shape,
                sub_network=self.use_pyramidbox,
                is_infer=True)
            infer_program, nmsed_out = network.infer(main_program)
            fetches = [nmsed_out]
            exe.run(startup_program)
            fluid.io.load_persistables(
                exe, self.model_dir, main_program=infer_program)

            self.exe, self.infer_program, self.fetches = exe, infer_program, fetches

    def detection_face(self, image):
        faceResult = ""
        shrink, max_shrink = self.get_shrink(image.size[1], image.size[0])
        try:
            det0 = self.detect_face(image, shrink)
        except Exception as e:
            return faceResult
        else:
            dets = det0
            keep_index = np.where(dets[:, 4] >= self.confs_threshold)[0]
            dets = dets[keep_index, :]  # all prediction included  in dets
            # 如果bbox数量为0，那么根据是否检测出车牌来决定以哪种方式输出到指定路径
            if len(dets) != 0:
                bboxes = dets[:, 0:4]
                faceResult = 'Face '
                for dex_b in range(len(bboxes)):
                    xmin, ymin, xmax, ymax = bboxes[dex_b]
                    xmin, ymin, xmax, ymax = math.ceil(xmin), math.ceil(ymin), math.floor(xmax), math.floor(ymax)
                    faceResult = faceResult + f'{str(xmin)} {str(ymin)} {str(xmax)} {str(ymax)} '
                    # print(xmin,ymin,xmax,ymax)
                    # print("this is shape",cv_image.shape)
        return faceResult

    def detect_face_image_video(self):
        video_capture = cv2.VideoCapture('E:/Video/Face.mp4')
        while True:
            ret, frame = video_capture.read()  # frame shape 640*480*3
            image = Image.fromarray(frame[..., ::-1])  # bgr to rgb
            shrink, max_shrink = self.get_shrink(image.size[1], image.size[0])
            try:
                det0 = self.detect_face(image, shrink)
            except Exception as e:
                raise
            else:
                dets = det0
                keep_index = np.where(dets[:, 4] >= self.confs_threshold)[0]
                dets = dets[keep_index, :]  # all prediction included  in dets
                # 如果bbox数量为0，那么根据是否检测出车牌来决定以哪种方式输出到指定路径
                if len(dets) != 0:
                    bboxes = dets[:, 0:4]
                    for dex_b in range(len(bboxes)):
                        xmin, ymin, xmax, ymax = bboxes[dex_b]
                        xmin, ymin, xmax, ymax = math.ceil(xmin), math.ceil(ymin), math.floor(xmax), math.floor(ymax)
                        cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (255, 0, 0), 2)
                        # print(xmin,ymin,xmax,ymax)
                        # print("this is shape",cv_image.shape)
            cv2.imshow('', frame)
            # Press Q to stop!
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        cv2.destroyAllWindows()

    def get_shrink(self, height, width):
        """
        Args:
            height (int): image height.
            width (int): image width.
        """
        # avoid out of memory
        max_shrink_v1 = (0x7fffffff / 577.0 / (height * width)) ** 0.5
        max_shrink_v2 = ((678 * 1024 * 2.0 * 2.0) / (height * width)) ** 0.5

        def get_round(x, loc):
            str_x = str(x)
            if '.' in str_x:
                str_before, str_after = str_x.split('.')
                len_after = len(str_after)
                if len_after >= 3:
                    str_final = str_before + '.' + str_after[0:loc]
                    return float(str_final)
                else:
                    return x

        max_shrink = get_round(min(max_shrink_v1, max_shrink_v2), 2) - 0.3
        if max_shrink >= 1.5 and max_shrink < 2:
            max_shrink = max_shrink - 0.1
        elif max_shrink >= 2 and max_shrink < 3:
            max_shrink = max_shrink - 0.2
        elif max_shrink >= 3 and max_shrink < 4:
            max_shrink = max_shrink - 0.3
        elif max_shrink >= 4 and max_shrink < 5:
            max_shrink = max_shrink - 0.4
        elif max_shrink >= 5:
            max_shrink = max_shrink - 0.5

        shrink = max_shrink if max_shrink < 1 else 1
        return shrink, max_shrink

    def detect_face(self, image, shrink):
        image_shape = [3, image.size[1], image.size[0]]
        if shrink != 1:
            h, w = int(image_shape[1] * shrink), int(image_shape[2] * shrink)
            image = image.resize((w, h), Image.ANTIALIAS)
            image_shape = [3, h, w]

        img = np.array(image)

        img = reader.to_chw_bgr(img)
        mean = [104., 117., 123.]
        scale = 0.007843
        img = img.astype('float32')
        img -= np.array(mean)[:, np.newaxis, np.newaxis].astype('float32')
        img = img * scale
        img = [img]
        img = np.array(img)

        detection, = self.exe.run(self.infer_program,
                                  feed={'image': img},
                                  fetch_list=self.fetches,
                                  return_numpy=False)
        detection = np.array(detection)
        # layout: xmin, ymin, xmax. ymax, score
        if np.prod(detection.shape) == 1:
            # print("No face detected")
            return np.array([[0, 0, 0, 0, 0]])
        det_conf = detection[:, 1]
        det_xmin = image_shape[2] * detection[:, 2] / shrink
        det_ymin = image_shape[1] * detection[:, 3] / shrink
        det_xmax = image_shape[2] * detection[:, 4] / shrink
        det_ymax = image_shape[1] * detection[:, 5] / shrink

        det = np.column_stack((det_xmin, det_ymin, det_xmax, det_ymax, det_conf))

        return det


if __name__ == '__main__':
    faceDetection = FaceDetection()
    faceDetection.detect_face_image_video()
    # addMosaic.detect_face_image(addMosaic.Prepocess_Decode("E:/AI/人车图片/多媒体/110102/行政区域/大栅栏街道12.jpg"))
